//
//  NetworkAnalysis.m
//  JP
//
//  Created by LiaoChenliang on 2022/4/13.
//  Copyright © 2022 LiaoChenliang. All rights reserved.
//

#import "AGNetworkAnalysis.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>
#include <arpa/inet.h>
#include <sys/sockio.h>
#include <net/ethernet.h>
#include <net/if.h>
#include <errno.h>
#include <net/if_dl.h>

#include "fishhook.h"
#import "DDLogWrapper.h"
#import <mach/mach_time.h>

@interface AGNetworkAnalysis()

@property (nonatomic, assign, readwrite) BOOL isStart;
@property (nonatomic, assign, readwrite) BOOL isEnableConsole;
@property (nonatomic, assign, readwrite) BOOL isEnableFile;
@property (nonatomic, assign, readwrite) AGNetworkAnalysisSection analysisSection;
@property (nonatomic, assign, readwrite) NSInteger platformType;
@property (nonatomic, copy, readwrite) NSString *logPath;
@property (nonatomic, strong) NSMutableDictionary *sockfdMap;
@property (nonatomic, assign) NSInteger currentRecvID;
@property (nonatomic, assign) NSInteger currentSendID;
@end

@implementation AGNetworkAnalysis

static AGNetworkAnalysis *__one;

+ (AGNetworkAnalysis *)shared {
static dispatch_once_t oneToken;
    dispatch_once(&oneToken, ^{
        __one = [[AGNetworkAnalysis alloc]init];

    });
    return __one;
}


static int     (*o_accept)(int, struct sockaddr * __restrict, socklen_t * __restrict);
static int     (*o_connect)(int, const struct sockaddr *, socklen_t);
static int     (*o_socket)(int, int, int);

static ssize_t (*o_recv)(int, void *, size_t, int);
static ssize_t (*o_recvfrom)(int, void *, size_t, int, struct sockaddr * __restrict, socklen_t * __restrict);
static ssize_t (*o_recvmsg)(int, struct msghdr *, int);
static ssize_t (*o_send)(int, const void *, size_t, int);
static ssize_t (*o_sendto)(int, const void *, size_t, int, const struct sockaddr *, socklen_t);
static ssize_t (*o_sendmsg)(int, const struct msghdr *, int);



//int h_connect(int sockfd, const struct sockaddr *addr, socklen_t addrlen) {
//    NSLog(@"h_connect = %d",sockfd);
//    return o_connect(sockfd,addr,addrlen);
//}
//
//int h_accept(int sockfd, struct sockaddr * __restrict addr, socklen_t * __restrict addrlen) {
//    NSLog(@"h_accept");
//    return o_accept(sockfd,addr,addrlen);
//}
//
//int h_socket(int sockfd, int b, int c){
//    NSLog(@"o_socket");
//    return o_socket(sockfd,b,c);
//}
/**********************************************************************************/

ssize_t h_recv(int sockfd, void *buf, size_t len, int flags) {
    if (isNetworkAnalysis(2)) {
        recvLog(sockfd, len);
    }
    return o_recv(sockfd,buf,len,flags);
}

ssize_t h_recvfrom(int sockfd, void *restrict buf, size_t len, int flags,
                   struct sockaddr *restrict src_addr,
                   socklen_t *restrict addrlen) {
    char ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &(((struct sockaddr_in*)src_addr)->sin_addr), ip, INET_ADDRSTRLEN);
    NSString *v = [NSString stringWithFormat:@"%s",ip];
    NSString *k = [NSString stringWithFormat:@"%d",sockfd];
    [[AGNetworkAnalysis shared].sockfdMap setValue:v forKey:k];
    if (isNetworkAnalysis(2)) {
        recvfromLog(sockfd, src_addr, len);
    }
    return o_recvfrom(sockfd,buf,len,flags,src_addr,addrlen);
}

ssize_t h_recvmsg(int sockfd, struct msghdr *msg, int flags){
    if (isNetworkAnalysis(2)) {
        recvmsgLog(sockfd, msg);
    }
    return o_recvmsg(sockfd,msg,flags);
}

ssize_t h_send(int sockfd, void *buf, size_t len, int flags) {
    if (isNetworkAnalysis(1)) {
        sendLog(sockfd, len);
    }
    return o_send(sockfd,buf,len,flags);
}

ssize_t h_sendto(int sockfd, const void *buf, size_t len, int flags,
            struct sockaddr *dest_addr, socklen_t addrlen){
    if (isNetworkAnalysis(1)) {
        sendtoLog(sockfd, dest_addr, len);
    }
    return o_sendto(sockfd,buf,len,flags,dest_addr,addrlen);
}

ssize_t h_sendmsg(int sockfd, struct msghdr *msg, int flags) {
    if (isNetworkAnalysis(1)) {
        sendmsgLog(sockfd, msg);
    }
    return o_sendmsg(sockfd,msg,flags);
}

void sendmsgLog(int sockfd, const struct msghdr * msg){
    uint64_t begin = mach_absolute_time();
    
    int msg_iovlen = msg->msg_iovlen;
    struct iovec *iov = msg->msg_iov;
    size_t len = 0;
    if (msg_iovlen>1) {
        struct iovec new_iov[msg_iovlen];
        for (int i = 0; i < msg_iovlen; i++) {
            new_iov[i] = iov[i];
            len = len + new_iov[i].iov_len;
        }
    }else{
        len = iov->iov_len;
    }
    const struct sockaddr * dest_addr = (((struct msghdr*)msg)->msg_name);
    if (dest_addr == NULL){return;}
    char ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &(((struct sockaddr_in*)dest_addr)->sin_addr), ip, INET_ADDRSTRLEN);
    NSString *ips = [NSString stringWithFormat:@"%s",ip];
    sys_sendtolog(ips,sockfd,len,begin);
}

void sendtoLog(int sockfd, struct sockaddr *dest_addr, size_t len){
    uint64_t begin = mach_absolute_time();
    
    char ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &(((struct sockaddr_in*)dest_addr)->sin_addr), ip, INET_ADDRSTRLEN);
    NSString *ips = [NSString stringWithFormat:@"%s",ip];
    sys_sendtolog(ips,sockfd,len,begin);
}

void sendLog(int sockfd, size_t len){
    uint64_t begin = mach_absolute_time();
    NSString *k = [NSString stringWithFormat:@"%d",sockfd];
    NSString *v = [[AGNetworkAnalysis shared].sockfdMap valueForKey:k];
    sys_sendtolog(v,sockfd,len,begin);
}

void sys_sendtolog(NSString* ip ,int sockfd, size_t len,uint64_t begin){
    
    NSString *logString = [[AGNetworkAnalysis shared].delegate networkAnalysisSend:[AGNetworkAnalysis shared].currentSendID remoteIP:ip contentLen:len];//[NSString stringWithFormat:@"send pt=%zd ip=%@ len=%zu",(long)[NetworkAnalysis shared].platformType,ip,len];
    [DDLogWrapper logInfo:logString];
    //NSLog(@"%@",logString);
    //uint64_t end = mach_absolute_time();
    //NSLog(@"t = %g ms", MachTimeToSecs(end - begin));
    //[DDLogWrapper logInfo:[NSString stringWithFormat:@"t = %g ms", MachTimeToSecs(end - begin)]];
    [AGNetworkAnalysis shared].currentSendID ++;
}

void recvmsgLog(int sockfd, const struct msghdr * msg){
    uint64_t begin = mach_absolute_time();
    
    int msg_iovlen = msg->msg_iovlen;
    struct iovec *iov = msg->msg_iov;
    size_t len = 0;
    if (msg_iovlen>1) {
        struct iovec new_iov[msg_iovlen];
        for (int i = 0; i < msg_iovlen; i++) {
            new_iov[i] = iov[i];
            len = len + new_iov[i].iov_len;
        }
    }else{
        len = iov->iov_len;
    }
    const struct sockaddr * dest_addr = (((struct msghdr*)msg)->msg_name);
    if (dest_addr == NULL){return;}
    char ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &(((struct sockaddr_in*)dest_addr)->sin_addr), ip, INET_ADDRSTRLEN);
    NSString *ips = [NSString stringWithFormat:@"%s",ip];
    sys_recvfromlog(ips,sockfd,len,begin);
}

void recvfromLog(int sockfd, struct sockaddr *dest_addr, size_t len){
    uint64_t begin = mach_absolute_time();
    
    char ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &(((struct sockaddr_in*)dest_addr)->sin_addr), ip, INET_ADDRSTRLEN);
    NSString *ips = [NSString stringWithFormat:@"%s",ip];
    sys_recvfromlog(ips,sockfd,len,begin);
}

void recvLog(int sockfd, size_t len){
    uint64_t begin = mach_absolute_time();
    NSString *k = [NSString stringWithFormat:@"%d",sockfd];
    NSString *v = [[AGNetworkAnalysis shared].sockfdMap valueForKey:k];
    sys_recvfromlog(v,sockfd,len,begin);
}

void sys_recvfromlog(NSString* ip ,int sockfd, size_t len,uint64_t begin){
    
    NSString *logString = [[AGNetworkAnalysis shared].delegate networkAnalysisRecv:[AGNetworkAnalysis shared].currentRecvID remoteIP:ip contentLen:len];//[NSString stringWithFormat:@"recv pt=%zd ip=%@ len=%zu",(long)[NetworkAnalysis shared].platformType,ip,len];
    [DDLogWrapper logInfo:logString];
    //NSLog(@"%@",logString);
    //uint64_t end = mach_absolute_time();
    //NSLog(@"t = %g ms", MachTimeToSecs(end - begin));
    //[DDLogWrapper logInfo:[NSString stringWithFormat:@"t = %g ms", MachTimeToSecs(end - begin)]];
    [AGNetworkAnalysis shared].currentRecvID ++;
}

bool isNetworkAnalysis(int section){
    if ([AGNetworkAnalysis shared].isStart) {
        if ([AGNetworkAnalysis shared].analysisSection == section || [AGNetworkAnalysis shared].analysisSection == AGNetworkAnalysisAll) {
            BOOL analysis =  [AGNetworkAnalysis shared].isEnableFile || [AGNetworkAnalysis shared].isEnableConsole;
            return analysis;
        }else{
            return false;
        }
    }else{
        return false;
    }
}

double MachTimeToSecs(uint64_t time){
    mach_timebase_info_data_t timebase;
    mach_timebase_info(&timebase);
    return (double)time * (double)timebase.numer /
    (double)timebase.denom /1e6;
}

- (void )start{
    self.isStart = true;
    if (self.isEnableFile) {
        NSString *paths = NSSearchPathForDirectoriesInDomains(NSDocumentDirectory, NSUserDomainMask, YES).firstObject;
        NSString *logsDirectory = [paths stringByAppendingPathComponent:@"NetworkAnalysisLogs"];
        self.logPath = logsDirectory;
        [DDLogWrapper addFileLogger:logsDirectory];
    }
    if (self.isEnableConsole){
        [DDLogWrapper addDDOSLogger];
    }
}

- (void )stop{
    self.isStart = false;
    [DDLogWrapper removeLogger];
}

- (void )setupAnalysisSection:(AGNetworkAnalysisSection)section{
    self.analysisSection = section;
    
    self.sockfdMap = [NSMutableDictionary dictionary];
    struct rebinding sh_recv = {
        "recv",
        h_recv,
        (void *)&o_recv

    };
    struct rebinding sh_recvfrom = {
        "recvfrom",
        h_recvfrom,
        (void *)&o_recvfrom

    };
    struct rebinding sh_recvmsg = {
        "recvmsg",
        h_recvmsg,
        (void *)&o_recvmsg

    };
    struct rebinding sh_send = {
        "send",
        h_send,
        (void *)&o_send

    };
    struct rebinding sh_sendmsg = {
        "sendmsg",
        h_sendmsg,
        (void *)&o_sendmsg

    };
    struct rebinding sh_sendto = {
        "sendto",
        h_sendto,
        (void *)&o_sendto

    };

    /*

     int h_accept(int a, struct sockaddr * __restrict b, socklen_t * __restrict c) {
         NSLog(@"h_accept");
         return o_accept(a,b,c);
     }

     int h_bind(int a, const struct sockaddr * b, socklen_t c){
         NSLog(@"h_bind");
         return o_bind(a,b,c);
     }

     int h_socket(int a, int b, int c){
         NSLog(@"o_socket");
         return o_socket(a,b,c);
     }
     */


//    struct rebinding sh_accept = {
//        "accept",
//        h_accept,
//        (void *)&o_accept
//
//    };
//    struct rebinding sh_connect = {
//        "connect",
//        h_connect,
//        (void *)&o_connect
//
//    };
//    struct rebinding sh_socket = {
//        "socket",
//        h_socket,
//        (void *)&o_socket
//
//    };
    rebind_symbols((struct rebinding[6]){
        sh_recv,
        sh_recvfrom,
        sh_recvmsg,
        sh_send,
        sh_sendmsg,
        sh_sendto,
//        sh_connect,
//        sh_accept,
//        sh_socket,
    }, 6);
}

- (void )setupPlatform:(NSInteger )type{
    self.platformType = type;
}
- (void )setupEnableConsole:(BOOL )enable{
    self.isEnableConsole = enable;
}
- (void )setupEnableFile:(BOOL )enable{
    self.isEnableFile = enable;
}
@end
